import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans

plt.rcParams['font.sans-serif'] = ['Simhei']
plt.rcParams['axes.unicode_minus'] = False

plt.figure()
x = np.loadtxt(r'../data/test.txt')
list1 = []
for i in range(10):
    k = i + 1
    model = KMeans(k)
    model.fit(x)
    list1.append(model.inertia_)
list2 = np.arange(1, 11)
plt.plot(list2, list1)

for i in range(10):
    plt.annotate(str(i + 1), xy=(list2[i], list1[i]))
plt.show()

plt.figure()
k = 4
model = KMeans(k)
model.fit(x)
carr = model.labels_
print(carr)
centers = model.cluster_centers_
print(centers)

plt.scatter(x[:, 0], x[:, 1], c=carr)
plt.scatter(centers[:, 0], centers[:, 1], c=['r', 'b', 'g', 'pink'], marker='*', s=180)
for i in range(4):
    plt.annotate('聚类 ' + str(i + 1), xy=(centers[i, 0], centers[i, 1]))

plt.show()
